package threshsig;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.math.BigInteger;
import java.security.GeneralSecurityException;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.interfaces.RSAPrivateKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPrivateKeySpec;
import java.util.Random;

/**
 * A Secret Key Share for an RSA (k,l) Threshold Scheme
 * 
 * Reference: "Practical Threshold Signatures",<br>
 * Victor Shoup (sho@zurich.ibm.com), IBM Research Paper RZ3121, 4/30/99<BR>
 * 
 * @author Steve Weis <sweis@mit.edu>
 */
public class KeyShare {
  // Constants and variables
  //............................................................................
  /** Secret key value */
  private BigInteger secret;
  /** Verifier used to authenticate self to other shares */
  private BigInteger verifier;
  private BigInteger groupVerifier;
  
  // TODO: This information should be moved to a GroupKey object
  // It's redundant to put into every share
  private BigInteger n;
  private final BigInteger delta;

  /** The secret key value used to sign messages. Should not be exposed */
  private BigInteger signVal;
  // TODO: Maybe get a better id scheme. Use verifier?
  private int id;

  // TODO-Rand: This is klunky.
  // Maybe use a utility class with a RNG for everyone
  private static SecureRandom random;
  private MessageDigest md;
  static {
    final byte[] randSeed = new byte[20];
    (new Random()).nextBytes(randSeed);
    random = new SecureRandom(randSeed);
  }
  
  // Hash-Scheme for RSA-compatible signatures
  private String rsaSigType = "SHA1withRSA";
  private RSAPrivateKey identityKey;

  
  // Constructors
  //............................................................................

  /**
   * Create a new share
   * 
   * @param id - the identifier of this share
   * @param secret - a secret value generated by a Dealer
   * @param n - the modulo of the group public key
   * @param delta - l! (group size factorial)
   */
  public KeyShare(final int id, final BigInteger secret, final BigInteger n, final BigInteger delta) {
    this.id = id;
    this.secret = secret;

    // Verifier must be set later using setVerifier()
    // TODO: Merge Dealer.generateShares and Dealer.generateVerifiers
    // and generate them simultaneously
    verifier = null;

    this.n = n;
    this.delta = delta;
    signVal = ThreshUtil.FOUR.multiply(delta).multiply(secret);
  }

  // Wrap/Unwrap Methods
  public static KeyShare unwrap(byte[] keydata) throws InvalidKeyException {
		try {
			ObjectInputStream oin = new ObjectInputStream(new ByteArrayInputStream(keydata));
			int id = oin.readInt();
			BigInteger secret = (BigInteger) oin.readObject();
			BigInteger n = (BigInteger) oin.readObject();
			BigInteger delta = (BigInteger) oin.readObject();
			KeyShare ks = new KeyShare(id, secret, n, delta);
		
			BigInteger verifier = (BigInteger)oin.readObject();
			BigInteger groupverifier = (BigInteger)oin.readObject();
			ks.setVerifiers(verifier, groupverifier);
			return ks;
		}
		catch(Exception ex) {
			throw new InvalidKeyException();
		}
	}
	
	public byte[] wrap() throws InvalidKeyException {
		ByteArrayOutputStream bo = new ByteArrayOutputStream();
		try {
			ObjectOutputStream oout = new ObjectOutputStream(bo);
			oout.writeInt(getId());
			oout.writeObject(getSecret());
			oout.writeObject(n);
			oout.writeObject(delta);
			oout.writeObject(verifier);
			oout.writeObject(groupVerifier);
			oout.close();
		} catch(Exception ex) {
			throw new InvalidKeyException();
		}
		return bo.toByteArray();
	}
  // Public Methods
  //............................................................................

  public int getId() {
    return id;
  }

  public BigInteger getSecret() {
    return secret;
  }
  
  public BigInteger getN() {
	  return n;
  }
  
  public BigInteger getDelta() {
	  return delta;
  }

  public void setVerifiers(final BigInteger verifier, final BigInteger groupVerifier) {
    this.verifier = verifier;
    this.groupVerifier = groupVerifier;
  }

  public BigInteger getVerifier() {
    return verifier;
  }
  
  public BigInteger getGroupVerifier() {
	  return groupVerifier;
  }

  public BigInteger getSignVal() {
    return signVal;
  }

  @Override
  public String toString() {
    return "KeyShare[" + id + "]";
  }

  /**
   * Create a SigShare and a Verifier for byte[] b<BR>
   * 
   * Refer to Shoup pg. 8 <BR>
   * 
   * @param b The array of bytes to produce a signature share for.
   * @return a sig share with a verifier
   */
  public SigShare sign(final byte[] b) {
    final BigInteger x = (new BigInteger(b)).mod(n);

    final int randbits = n.bitLength() + 3 * ThreshUtil.L1;

    // r \elt (0, 2^L(n)+3*l1)
    final BigInteger r = (new BigInteger(randbits, random));
    final BigInteger vprime = groupVerifier.modPow(r, n);
    final BigInteger xtilde = x.modPow(ThreshUtil.FOUR.multiply(delta), n);
    final BigInteger xprime = xtilde.modPow(r, n);

    BigInteger c = null;
    BigInteger z = null;
    // Try to generate C and Z
    try {
      md = MessageDigest.getInstance("SHA");
      md.reset();

      // debug("v: " + groupVerifier.mod(n));
      md.update(groupVerifier.mod(n).toByteArray());

      // debug("xtilde: " + xtilde);
      md.update(xtilde.toByteArray());

      // debug("vi: " + verifier.mod(n));
      md.update(verifier.mod(n).toByteArray());

      // debug("xi^2: " + x.modPow(signVal,n).modPow(TWO,n));
      md.update(x.modPow(signVal, n).modPow(ThreshUtil.TWO, n).toByteArray());

      // debug("v': "+ vprime);
      md.update(vprime.toByteArray());

      // debug("x': " + xprime);
      md.update(xprime.toByteArray());
      c = new BigInteger(md.digest()).mod(n);
      z = (c.multiply(secret)).add(r);
    } catch (final java.security.NoSuchAlgorithmException e) {
      debug("Provider could not locate SHA message digest .");
      e.printStackTrace();
    }

    final Verifier ver = new Verifier(z, c, verifier, groupVerifier);

    return new SigShare(id, x.modPow(signVal, n), ver);
  }

  /**
   * Set type of signature (hash and padding) for RSA compatibility.
   * There must be a RSA Signature provider installed for the
   * specified type
   * @param type 
   */
  public void setRSASigType(String type) {
	  rsaSigType = type;
  }
  
  /**
   * Calculate Hash and add padding, then create a SigShare
   * using SigShare.sign<BR>
   * 
   * @param b The array of bytes to produce a signature share for.
   * @return a sig share with a verifier
   * @throws NoSuchAlgorithmException 
   */
  public SigShare rsasign(final byte[] b) throws NoSuchAlgorithmException {
	  byte[] data = paddeddigest(b);
	  return sign(data);
  }
  
  /**
   * Calculate Hash and add padding compatible to standard RSA
   * signature.
   * Key size is used defined in the constructure, hash/padding type
   * can be set by setRSASigType and defaults to "SHA1withRSA"
   *  
   * @param b The array of bytes to produce a hash with padding for.
   * @return padded hash
   */
  public byte[] paddeddigest(final byte[] b) throws NoSuchAlgorithmException {
	  if(identityKey==null) {
		  // initialize identityKey upon first use
		  KeyFactory key_factory = KeyFactory.getInstance("RSA");
	  	  RSAPrivateKeySpec rsaspec = new RSAPrivateKeySpec(n, ThreshUtil.ONE);
	  	  try {
			  identityKey = (RSAPrivateKey) key_factory.generatePrivate( rsaspec );
		  } catch (InvalidKeySpecException e) {
			  throw new RuntimeException("FATAL ERROR: should never happen",e);
		  }
	  }
	  return paddeddigest(b, identityKey, rsaSigType);
  }
  
  /**
   * Calculate Hash and add padding compatible to standard RSA
   * signature.
   *  
   * @param b The array of bytes to produce a hash with padding for.
   * @param n The modulus of the RSA key (used for determining the key size)
   * @param rsaSigType String (as in JCE) to select Hash/Padding (e.g., "SHA1withRSA"),
   *                   which needs to be supported by some JCE crypto provider.
   * @return padded hash
   */
  public static byte[] paddeddigest(final byte[] b, BigInteger n, String rsaSigType) {
  	  try {
  		  KeyFactory key_factory = KeyFactory.getInstance("RSA");
  		  RSAPrivateKeySpec rsaspec = new RSAPrivateKeySpec(n, ThreshUtil.ONE);
		  RSAPrivateKey identityKey = (RSAPrivateKey) key_factory.generatePrivate( rsaspec );
		  return paddeddigest(b, identityKey, rsaSigType);
	  } catch (GeneralSecurityException e) {
		  throw new RuntimeException("FATAL ERROR: should never happen",e);
	  }
  }
  
  private static byte[] paddeddigest(final byte[] b, RSAPrivateKey identityKey, String rsaSigType) throws NoSuchAlgorithmException {
	  /*
	   * As JCA does not offer any public method for creating RSA padding,
	   * we simply create a regular RSA signature with private exponent 1
	   * the signing key is either created once on first use (instance method)
	   * or created on each call (static class method)
	   */
	  Signature instance = Signature.getInstance(rsaSigType);
	  byte[] paddedhash = null;
	  try {
		  instance.initSign(identityKey);
		  instance.update(b);
		  paddedhash = instance.sign();
	  } catch (GeneralSecurityException e) {
		  // should never ever happen
		  e.printStackTrace();
	  }
	  return paddedhash;
  }
  
  
  // Debugging
  //............................................................................
  private static void debug(final String s) {
    System.err.println("KeyShare: " + s);
  }
}
